#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from datetime import datetime


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def AssignJob(job_id: str, downloader: str, file_count: int, priority_field=""):
    client_joblist, collection_joblist = InitMongoDB("10.10.116.28", 27017, "ant_nest", "joblist")
    client_datalist, collection_datalist = InitMongoDB("10.10.116.28", 27017, "ant_nest", "datalist")

    count = 0

    if priority_field != "":
        for result in collection_datalist.find(
                {"type": priority_field, "data_store_path": {"$exists": 0}, "JobID": {"$exists": 0}}, {"_id": 1}):
            count += 1
            collection_datalist.update_one({"_id": result["_id"]},
                                           {"$set": {"download_way": "external", "JobID": job_id}})
            if count == file_count:
                break
    if count < file_count:
        for result in collection_datalist.find({"data_store_path": {"$exists": 0}, "JobID": {"$exists": 0}},
                                               {"_id": 1}):
            count += 1
            collection_datalist.update_one({"_id": result["_id"]},
                                           {"$set": {"download_way": "external", "JobID": job_id}})
            if count == file_count:
                break

    job_record = {}
    job_record["JobID"] = job_id
    job_record["job_create_date"] = datetime.today()
    job_record["downloader"] = downloader
    job_record["file_count"] = count
    job_record["status"] = "assigned"
    if count != 0:
        collection_joblist.insert_one(job_record)
        print("Job assigned " + job_id)

    CloseMongoDB(client_datalist)
    CloseMongoDB(client_joblist)


if __name__ == '__main__':
    # AssignJob("dlst000001", "Yuan Liyun", 2000, priority_field="human lung")
    # AssignJob("dlst000002", "Yuan Liyun", 2000, priority_field="human lung")
    # AssignJob("dlst000003", "Yuan Liyun", 2000, priority_field="cancer")
    # AssignJob("dlst000004", "Yuan Liyun", 2000, priority_field="human gut")
    # AssignJob("dlst000005", "Yuan Liyun", 2000, priority_field="human gut")
    # AssignJob("dlst000006", "Chen Yujie", 2000, priority_field="human and animal gut")
    # AssignJob("dlst000007", "Chen Yujie", 2000)
    # AssignJob("dlst000008", "Chen Yujie", 2000)
    # AssignJob("dlst000009", "Chen Yujie", 2000)
    # AssignJob("dlst000010", "Chen Yujie", 2000)
    # AssignJob("dlst000011", "Ling Yunchao", 2000)
    # AssignJob("dlst000012", "Ling Yunchao", 2000)
    # AssignJob("dlst000013", "Ling Yunchao", 2000)
    # AssignJob("dlst000014", "Ling Yunchao", 2000)
    # AssignJob("dlst000015", "Ling Yunchao", 2000)
    AssignJob("dlst000016", "Yuan Liyun", 2000)
    AssignJob("dlst000017", "Yuan Liyun", 2000)
    AssignJob("dlst000018", "Yuan Liyun", 2000)
    AssignJob("dlst000019", "Yuan Liyun", 2000)
    AssignJob("dlst000020", "Yuan Liyun", 2000)
    AssignJob("dlst000021", "Chen Yujie", 2000)
    AssignJob("dlst000022", "Chen Yujie", 2000)
    AssignJob("dlst000023", "Chen Yujie", 2000)
    AssignJob("dlst000024", "Chen Yujie", 2000)
    AssignJob("dlst000025", "Chen Yujie", 2000)
    AssignJob("dlst000026", "Ling Yunchao", 2000)
    AssignJob("dlst000027", "Ling Yunchao", 2000)
    AssignJob("dlst000028", "Ling Yunchao", 2000)
    AssignJob("dlst000029", "Ling Yunchao", 2000)
    AssignJob("dlst000030", "Ling Yunchao", 2000)
